from scipy.spatial.transform import Rotation as R
import os
import json, tqdm

def get_cam_para(para_path):
    
    calibrated_sensor_dict = {}
    
    for cam in tqdm.tqdm(os.listdir(para_path), desc='Getting camera para...'):
        if cam.startswith('CAM'):
            with open(os.path.join('param' , cam)) as f:
                cam_data = json.load(f)
            intrinsic_matrix = cam_data['intrinsic_param']['camera_matrix']
            axis_angle = cam_data['extrinsic_param']['rotation']
            translation = cam_data['extrinsic_param']['translation']
            quat = R.from_rotvec(axis_angle).as_quat().tolist()
            calibrated_sensor_dict[cam.split('.')[0]] = {
                'intrinsics' : intrinsic_matrix,
                'quat' : [quat[3], quat[0], quat[1], quat[2]],
                'translation' : translation
            }
        elif cam.startswith('LIDAR'):
            with open(os.path.join('param' , cam)) as f:
                cam_data = json.load(f)
            intrinsic_matrix = cam_data['intrinsic_param']
            axis_angle = cam_data['extrinsic_param']['rotation']
            translation = cam_data['extrinsic_param']['translation']
            quat = R.from_rotvec(axis_angle).as_quat().tolist()
            calibrated_sensor_dict[cam.split('.')[0]] = {
                'intrinsics' : intrinsic_matrix,
                'quat' : [quat[3], quat[0], quat[1], quat[2]],
                'translation' : translation
            }

    return calibrated_sensor_dict
